/**
 * 向量搜索工具
 * 
 * 基于向量相似度搜索知识库中的相关内容
 */

import { prisma } from "./prisma";
import { createEmbeddings, calculateCosineSimilarity } from "./embeddings";
import { SearchResult } from "@/types/chat";
import { Vector, Document } from "@prisma/client";

// 扩展Vector类型，包含document关系
type VectorWithDocument = Vector & {
  document: {
    id: string;
    name: string;
  }
};

/**
 * 在知识库中搜索相关内容
 * @param query 搜索查询文本
 * @param knowledgeBaseId 知识库ID
 * @param limit 最大返回结果数
 * @param similarityThreshold 相似度阈值，低于此值的结果将被过滤
 * @returns 搜索结果数组
 */
export async function searchKnowledgeBase(
  query: string,
  knowledgeBaseId: string,
  limit: number = 5,
  similarityThreshold: number = 0.7
): Promise<SearchResult[]> {
  try {
    // 验证参数
    if (!query || query.trim() === '') {
      return [];
    }
    
    if (!knowledgeBaseId) {
      throw new Error("知识库ID不能为空");
    }
    
    // 生成查询的嵌入向量
    const queryEmbedding = await createEmbeddings(query);
    
    if (!queryEmbedding || queryEmbedding.length === 0) {
      throw new Error("无法为查询生成嵌入向量");
    }
    
    // 获取知识库中的所有向量
    const vectors = await prisma.vector.findMany({
      where: {
        document: {
          knowledgeBaseId: knowledgeBaseId
        }
      },
      include: {
        document: {
          select: {
            id: true,
            name: true
          }
        }
      }
    });
    
    if (vectors.length === 0) {
      return [];
    }
    
    // 计算相似度并排序
    const results = vectors
      .map((vector: VectorWithDocument) => {
        // 计算向量相似度
        const similarity = calculateCosineSimilarity(
          queryEmbedding,
          vector.embedding as unknown as number[]
        );
        
        return {
          documentId: vector.documentId,
          documentName: vector.document.name,
          content: vector.content,
          similarity
        };
      })
      // 按相似度降序排序
      .sort((a: SearchResult, b: SearchResult) => b.similarity - a.similarity)
      // 过滤掉相似度低于阈值的结果
      .filter((result: SearchResult) => result.similarity >= similarityThreshold)
      // 限制结果数量
      .slice(0, limit);
    
    return results;
  } catch (error) {
    console.error("知识库搜索错误:", error);
    throw new Error(`搜索知识库失败: ${(error as Error).message}`);
  }
}

// 扩展Document类型，包含vectors关系
type DocumentWithVectors = Document & {
  vectors: Vector[];
};

/**
 * 执行语义搜索（使用关键词而非向量）
 * 当向量搜索不可用或需要备用方案时使用
 * @param query 搜索关键词
 * @param knowledgeBaseId 知识库ID
 * @param limit 最大返回结果数
 * @returns 搜索结果数组
 */
export async function semanticSearch(
  query: string,
  knowledgeBaseId: string,
  limit: number = 5
): Promise<SearchResult[]> {
  try {
    // 验证参数
    if (!query || query.trim() === '') {
      return [];
    }
    
    if (!knowledgeBaseId) {
      throw new Error("知识库ID不能为空");
    }
    
    // 准备搜索关键词
    const keywords = query
      .toLowerCase()
      .split(/\s+/)
      .filter(word => word.length > 2);
    
    if (keywords.length === 0) {
      return [];
    }
    
    // 获取知识库中的所有文档
    const documents = await prisma.document.findMany({
      where: {
        knowledgeBaseId
      },
      include: {
        vectors: true
      }
    });
    
    if (documents.length === 0) {
      return [];
    }
    
    // 对每个向量执行关键词匹配
    const results: SearchResult[] = [];
    
    for (const doc of documents as DocumentWithVectors[]) {
      for (const vector of doc.vectors) {
        // 简单计算关键词匹配分数
        let matchScore = 0;
        const contentLower = vector.content.toLowerCase();
        
        for (const keyword of keywords) {
          if (contentLower.includes(keyword)) {
            matchScore += 1;
          }
        }
        
        // 归一化分数
        const similarity = keywords.length > 0 ? matchScore / keywords.length : 0;
        
        // 只添加有匹配的结果
        if (similarity > 0) {
          results.push({
            documentId: doc.id,
            documentName: doc.name,
            content: vector.content,
            similarity
          });
        }
      }
    }
    
    // 按相似度分数排序并限制结果数
    return results
      .sort((a: SearchResult, b: SearchResult) => b.similarity - a.similarity)
      .slice(0, limit);
  } catch (error) {
    console.error("语义搜索错误:", error);
    return [];
  }
} 